#include "../../include/core/container.h"
#include <cstring>
#include <iostream>

#define INPUT_SIZE 16
#define OUTPUT_SIZE 16
#define KEY_SIZE 16

extern "C" {
void AES_OP(MemRef<unsigned char, 1> *inputs, int msgLen,
            MemRef<unsigned char, 1> *outputs, MemRef<unsigned char, 1> *key,
            MemRef<unsigned char, 1> *sBoxAes);
}

uint8_t exampleInput[INPUT_SIZE] = {'0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7'};
uint8_t exampleKey[KEY_SIZE] = {'A', '0', 'B', '1', 'C', '2', 'D', '3','E', '4', 'F', '5', 'G', '6', 'H', '7'};
uint8_t sBox = {'0'};

MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> key{exampleKey, KEY_SIZE};
uint8_t *exampleOutput = new uint8_t[OUTPUT_SIZE];
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};
MemRef<unsigned char, 1> sBoxAes{sBox, 1};

void printResult(uint8_t *output) {
  // Print the output.
  std::cout << "[ ";
  for(size_t i = 0; i < OUTPUT_SIZE; i +=4) {
    uint32_t tempNum;
    memcpy(&tempNum, output +i, sizeof(uint32_t));
    printf("%x", tempNum);
  }
  std::cout << " ]" << std::endl;
}

void printInput(uint8_t *input) {
  // Print the input.
  std::cout << "Input : [";
  for (size_t i = 0; i < INPUT_SIZE; i++) {
    std::cout << input[i];
  }
  std::cout << "] ==> ";
}

int main() {
  int keyLen = 16;
  AES_OP(&input, INPUT_SIZE, &output, &key, &sBoxAes);
  printInput(exampleInput);
  printResult(exampleOutput);

  std::cout << "Key " << ": [";
  for (size_t i = 0; i < keyLen; i++) {
    std::cout << key[i];
  }
  std::cout << "]" << std::endl;

  return 0;
}
